[None][test] Declare supported attention test phases#16718
Conversation
yihwang-nv
changed the title
760ef7e to
9543131
Compare
yihwang-nv
changed the title
WalkthroughAttention model configurations now declare optional context or generation phases. Backend attention test expansion selects configured phases or all available phases by default, replacing MLA-specific context branching with phase-based selection. ChangesAttention phase selection
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
tests/unittest/_torch/attention/test_attention_backends.py (1)
230-232: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate the new test function.
Add precise parameter and
-> Noneannotations, consistent with the Python typing guideline.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/test_attention_backends.py` around lines 230 - 232, Update test_phases_to_run with precise type annotations for the declared and expected parameters, and add a -> None return annotation. Use the existing typing conventions and types implied by ModelAttnConfig phases and _phases_to_run.Source: Coding guidelines
tests/unittest/_torch/attention/model_attn_config.py (1)
77-77: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse PEP 604 syntax for the new optional field.
- phases: Optional[tuple[AttentionPhase, ...]] = None + phases: tuple[AttentionPhase, ...] | None = NoneAs per coding guidelines, prefer built-in generic types and
|.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/model_attn_config.py` at line 77, Update the phases field annotation to use PEP 604 union syntax with the built-in tuple generic, replacing Optional while preserving the existing AttentionPhase tuple and None default.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/vanilla.py (2)
589-591: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
cache_indices = [list(block_ids) for block_ids in metadata.block_ids_per_seq]is duplicated verbatim between_mla_forward_generationandforward.Both the MLA-generation path (lines 589-591) and the standard
forward()path (lines 731-733) materialize per-request block-id lists identically. Extracting this into a small metadata/helper accessor would reduce the risk of the two call sites diverging as the paged-cache logic evolves further.Also applies to: 731-733
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/vanilla.py` around lines 589 - 591, Extract the duplicated block-ID materialization into a shared metadata helper or accessor, then update both `_mla_forward_generation` and `forward` to use it instead of constructing `cache_indices` inline. Preserve the existing per-request list output and avoid changing the paged-cache behavior.
204-219: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winWrite path silently trusts that new-token
block_idsare neverBAD_PAGE_INDEX.Unlike the read path (
_gather_paged_kv/_gather_paged_mla_latent), which explicitly checksblk == BAD_PAGE_INDEX, this write loop doeskv_cache_tensor[blk, 0]unconditionally. SinceBAD_PAGE_INDEX == -1, if this invariant is ever violated (e.g. a KV-manager bug hands back an evicted/unallocated block for a new-token position), PyTorch's negative indexing silently wraps to the last page instead of raising, corrupting an unrelated request's cache with no error signal. An assertion here would convert a silent-corruption bug into a loud, debuggable failure.🛡️ Proposed defensive assertion
while written < kv_len: pos = past_seen_token + written blk = block_ids[pos // tokens_per_block] + assert blk != BAD_PAGE_INDEX, ( + "New KV token position mapped to an invalid page; " + "cache manager must allocate blocks before writing new tokens.") off = pos % tokens_per_block🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/vanilla.py` around lines 204 - 219, Add a defensive assertion in the k/v cache write loop before indexing kv_cache_tensor, validating that blk is not BAD_PAGE_INDEX. Keep the existing block and offset calculations unchanged, and ensure an invalid new-token block fails loudly instead of allowing negative indexing to write to the final page.tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py (1)
94-102: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winRegex regression guard is narrower than its sibling and may miss reintroductions of the bug it guards against.
test_flash_mla_tokens_per_block_propagates_to_kv_cache_config(lines 61-91, unchanged) captures the entire branch body via((?:[ \t]+.*\n)+)and then searches within it. This new test instead requires the offending assignment to appear as the immediate next line after theif(\s*tokens_per_block\s*=\s*max_num_tokensright after the\n). If the anti-pattern were reintroduced with an intervening comment/blank line, inside anelif, or a few lines further into the branch body, this regex would not match and the test would pass while VANILLA is still (incorrectly) linearized.Consider mirroring the sibling test's approach: capture the full branch body first, then assert the offending pattern is absent from it.
♻️ Suggested widening
def test_vanilla_preserves_configured_tokens_per_block(): """Vanilla supports paged KV and must honor the configured page size.""" source = _get_create_py_executor_source() - assert not re.search( - r"if\s+llm_args\.attn_backend\s*==\s*[\"']VANILLA[\"']\s*:\s*\n" - r"\s*tokens_per_block\s*=\s*max_num_tokens", - source, - ), "Vanilla must not linearize the KV cache by overriding tokens_per_block." + vanilla_block_match = re.search( + r"if\s+llm_args\.attn_backend\s*==\s*[\"']VANILLA[\"']\s*:\s*\n" + r"((?:[ \t]+.*\n)+)", + source, + ) + if vanilla_block_match: + assert "tokens_per_block = max_num_tokens" not in vanilla_block_match.group(1), ( + "Vanilla must not linearize the KV cache by overriding tokens_per_block.")Test coverage summary (QA). Added test function:
test_vanilla_preserves_configured_tokens_per_block— new, source-inspection regression guard ensuring VANILLA doesn't linearize the KV cache viatokens_per_block = max_num_tokens. No existing tests modified/removed. Coverage verdict: needs follow-up — registration intests/integration/test_lists/test-db/orqa/could not be confirmed from provided context.
As per path instructions fortests/**: "Always produce a test coverage summary, even if no issues are found."#!/bin/bash rg -n "test_py_executor_creator_flash_mla_tokens_per_block" tests/integration/test_lists --glob '*.txt' --glob '*.yml' -C2🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py` around lines 94 - 102, The regex in test_vanilla_preserves_configured_tokens_per_block is too narrow and only detects an immediate assignment after the VANILLA condition. Mirror test_flash_mla_tokens_per_block_propagates_to_kv_cache_config by capturing the entire VANILLA branch body, then assert that body does not contain tokens_per_block = max_num_tokens, including when separated by comments, blank lines, or branch statements.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/vanilla.py`:
- Around line 164-183: Update _gather_paged_mla_latent to return an
appropriately shaped empty tensor when kv_len is less than or equal to zero,
matching the existing guard behavior in _gather_paged_kv; keep the current paged
gathering logic unchanged for positive lengths.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/vanilla.py`:
- Around line 589-591: Extract the duplicated block-ID materialization into a
shared metadata helper or accessor, then update both `_mla_forward_generation`
and `forward` to use it instead of constructing `cache_indices` inline. Preserve
the existing per-request list output and avoid changing the paged-cache
behavior.
- Around line 204-219: Add a defensive assertion in the k/v cache write loop
before indexing kv_cache_tensor, validating that blk is not BAD_PAGE_INDEX. Keep
the existing block and offset calculations unchanged, and ensure an invalid
new-token block fails loudly instead of allowing negative indexing to write to
the final page.
In `@tests/unittest/_torch/attention/model_attn_config.py`:
- Line 77: Update the phases field annotation to use PEP 604 union syntax with
the built-in tuple generic, replacing Optional while preserving the existing
AttentionPhase tuple and None default.
In `@tests/unittest/_torch/attention/test_attention_backends.py`:
- Around line 230-232: Update test_phases_to_run with precise type annotations
for the declared and expected parameters, and add a -> None return annotation.
Use the existing typing conventions and types implied by ModelAttnConfig phases
and _phases_to_run.
In
`@tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py`:
- Around line 94-102: The regex in
test_vanilla_preserves_configured_tokens_per_block is too narrow and only
detects an immediate assignment after the VANILLA condition. Mirror
test_flash_mla_tokens_per_block_propagates_to_kv_cache_config by capturing the
entire VANILLA branch body, then assert that body does not contain
tokens_per_block = max_num_tokens, including when separated by comments, blank
lines, or branch statements.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5f9e76bb-bd51-4463-98ea-109b8a9b9883
📒 Files selected for processing (7)
tensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytests/unittest/_torch/attention/backend_capability.pytests/unittest/_torch/attention/model_attn_config.pytests/unittest/_torch/attention/test_attention_backends.pytests/unittest/_torch/attention/test_vanilla_attention.pytests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
|
/bot run --disable-fail-fast |
|
PR_Github #60952 [ run ] triggered by Bot. Commit: |
|
PR_Github #60952 [ run ] completed with state
|
Replace mla_context with phases so each model attention config declares whether it supports context, generation, or both. Keep the existing context, generation, and mixed cases for standard attention when phases is None. Explicit MLA phases select only their declared atomic paths, so a config supporting context and generation does not imply an MLA mixed batch. Use immutable typed tuples for phase declarations. Signed-off-by: Yihan Wang <yihwang@nvidia.com>
9543131 to
f1abeba
Compare
|
/bot run |
|
PR_Github #61200 [ run ] triggered by Bot. Commit: |
|
PR_Github #61200 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61258 [ run ] triggered by Bot. Commit: |
|
PR_Github #61258 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61487 [ run ] triggered by Bot. Commit: |
|
PR_Github #61487 [ run ] completed with state |
3 participants
Summary
Testing
Dev Engineer Review
mla_contextwith an explicit, immutablephases: tuple[AttentionPhase, ...] | NonewhereAttentionPhase = Literal["ctx", "gen"].phases=("gen",)(for_mlavariants) andphases=("ctx",)(for*_mla_ctxvariants)._phases_to_run(cfg, available_phases)and uses it to generate backend cases based oncfg.phases(defaulting to all available phases when unspecified).{ctx, dec, mix}to{ctx, gen, mix}.QA Engineer Review
tests/): modifiedtests/unittest/_torch/attention/model_attn_config.pytests/unittest/_torch/attention/test_attention_backends.pytest_attention_backends.py(the existingtest_attention_backendandtest_split_consistencyremain); changes are to helper logic and parametrization generation.tests/integration/test_lists/:tests/integration/test_lists/test-db/l0_*.ymlincludesunittest/_torch/attention(covers this module).tests/integration/test_lists/waives.txtcontains SKIPs for specifictest_attention_backend[...]parametrizations in this module.